#############################################################
# This file reads pressure data for Brazilian Congress 
#   and estimate Positive-Trait Item Response Model
############################################################

rm(list=ls())

print(version)

# platform       x86_64-apple-darwin17.0     
# arch           x86_64                      
# os             darwin17.0                  
# system         x86_64, darwin17.0          
# status                                     
# major          4                           
# minor          2.1                         
# year           2022                        
# month          06                          
# day            23                          
# svn rev        82513                       
# language       R                           
# version.string R version 4.2.1 (2022-06-23)
# nickname       Funny-Looking Kid           




library (openxlsx)
library (car)
library (data.table)
library (gtools)
library (dplyr)
library (runjags)
library (coda)
library (random)

#### Call data ####
savePathRep <- "/all_data/"

PT.Pressure<- read.table (file=paste0(savePathRep, "/PT_pressure_analysis.txt"), sep="\t", header=TRUE)
PMDB.Pressure<- read.table (file=paste0(savePathRep, "/PMDB_pressure_analysis.txt"), sep="\t", header=TRUE)
PSDB.Pressure<- read.table (file=paste0(savePathRep, "/PSDB_pressure_analysis.txt"), sep="\t", header=TRUE)
DEM.Pressure<- read.table (file=paste0(savePathRep, "/DEM_pressure_analysis.txt"), sep="\t", header=TRUE)
PP.Pressure<- read.table (file=paste0(savePathRep, "/PP_pressure_analysis.txt"), sep="\t", header=TRUE)
PR.Pressure<- read.table (file=paste0(savePathRep, "/PR_pressure_analysis.txt"), sep="\t", header=TRUE)

tmp <- list (PT.Pressure, PMDB.Pressure, PSDB.Pressure
                       , DEM.Pressure, PP.Pressure, PR.Pressure)
y.count<- array(unlist (tmp), dim=c(nrow(tmp[[1]]), ncol(tmp[[1]]), 6) )


# Number of anchors (for identification)
A <- 3
# Number of bridges (to place pressure on similar scale)
B <- 2
 

##Random Seed
rn <- 10082

# Write model
jags.PTIRT <- "model{
## Likelihood
for (i in 1:N){
  for (j in 1:M){
    for (p in 1:P){
      y[i,j,p] ~ dbern(pi[i,j,p])
      probit(pi[i,j,p]) <- beta[j,p]*log(X[i,p]) - log(alpha[j,p])
    }
  }
}
## Priors

for (j in 1:B) {  # B bridge items should appear first
  beta.c[j] ~ dnorm(0, 0.25)
  for (p in 1:P){
    beta[j,p] <- beta.c[j]
  }
}

   # A anchor items should appear next
  for (p in 1:P){
beta[3,p]  ~ dnorm(0, 0.25)I(,-.3)
}
 for (p in 1:P){
beta[4,p]  ~ dnorm(0, 0.25)I(.3,)
}
 
for (p in 1:P){
beta[5,p]  ~ dnorm(0, 0.25)I(0,)
}

for (j in (B+A+1):M) {   # M-B-A regular items appear last
  for (p in 1:P) {
    beta[j,p] ~ dnorm(0, 0.25)
  }
}
for (j in 1:M) {
  for (p in 1:P) {
    alpha[j,p] ~ dgamma(.1,.1) 
  }
}
for (i in 1:N) {
  for (p in 1:P) {
    X[i,p] ~ dlnorm(0,10)
  }
}
}"

# Gather data
PTIRT.data <- dump.format(list(y= y.count
                               , N=dim(y.count)[1]
                               , M=dim(y.count)[2]
                               , P=dim(y.count)[3]
                               , A=A
                               , B=B
))

# Gather parameters
PTIRT.parameters = c("X", "alpha", "beta", "deviance")

# Initialize parameters
PTIRT.inits.1 <- function() {
  dump.format(
    list(
      X = matrix (data = c(rlnorm(dim(y.count)[1]*dim(y.count)[3], 0, 1))
                  , ncol=dim(y.count)[3]
                  , byrow=T)
      , alpha = matrix (data = rep(rgamma(dim(y.count)[2], .1, .1), 6),
                        ncol = 6)
      , beta = matrix (data = c(runif (B*6, -1, 1),
                                rep(-1, 6),rep(1, 6),rep(1, 6),
                                rep(rnorm (dim(y.count)[2]-A-B), 6)),
                       ncol = 6, byrow = T)
      , beta.c=rnorm(B)
      ,.RNG.name="base::Wichmann-Hill"
      ,.RNG.seed=rn
    )
  )
}


# Initialize parameters
PTIRT.inits.2 <- function() {
  dump.format(
    list(
      X =  matrix (data = c(rlnorm(dim(y.count)[1]*dim(y.count)[3], 0, 1))
                   , ncol=dim(y.count)[3]
                   , byrow=T)
      , alpha = matrix (data = rep(rgamma(dim(y.count)[2], .1, .1), 6),
                        ncol = 6)
      , beta = matrix (data = c(runif (B*6, -1, 1),
                                rep(-1, 6),rep(1, 6),rep(1, 6),
                                rep(rnorm (dim(y.count)[2]-A-B), 6)),
                       ncol = 6, byrow = T)
      , beta.c=rnorm(B)
      ,.RNG.name="base::Wichmann-Hill"
      ,.RNG.seed=rn+1
    )
  )
}

# run model
set.seed(rn+2)
PTIRT.model <- run.jags(
  model=jags.PTIRT,
  monitor=PTIRT.parameters,
  method="parallel",
  n.chains=2,
  data=PTIRT.data,
  inits=list (PTIRT.inits.1(), PTIRT.inits.2()),
  thin=30, burnin=5000, sample=200,
  check.conv=FALSE, plots=FALSE)

# Gather MCMC chains
chainsPTIRT <- mcmc.list(list (PTIRT.model$mcmc[[1]], PTIRT.model$mcmc[[2]]))

# Check convergence
gD <- gelman.diag (chainsPTIRT, multivariate=F) 
summary (gD[[1]][,1]) 

### Uncomment the next line to save "brazil_ptirt_mcmc.Rda"
# save (chainsPTIRT, file="brazil_ptirt_mcmc.Rda")

